package com.dap.servicegovernance.trace;

import com.dap.commons.Constants;
import com.dap.core.context.DapThreadContext;
import com.dap.servicegovernance.constants.TraceConstants;
import com.dap.utils.LongStringIdUtil;
import com.dap.utils.TraceIdUtils;
import org.slf4j.MDC;
import org.springframework.util.StringUtils;

import java.util.Stack;

/**
 * trace切面处理类
 */
public class TraceAOPRpc extends TraceAOPLocal {
    /**
     * 构造函数
     * @param logLevel 打印的日志级别
     */
    public TraceAOPRpc(String logLevel) {
        super(logLevel);
    }

    /**
     * 方法调用前处理
     */
    @Override
    protected String beforeMethodInvoked(String className, String methodName, Object[] args) {
        manageRpcIdBeforeMethodInvoked(className);

        TraceIdUtils.generateTraceId(false);

        logBeforeMethodInvoked(className, methodName, args);

        return DapThreadContext.getContext(TraceConstants.RPC_PARENT_CLASS_NAME);
    }
    /**
     * 方法调用后处理
     */
    @Override
    protected void afterMethodInvoked(String className, String methodName, long beginTime, Object invokeResult, String parentClassName) {
        logAfterMethodInvoked(className, methodName, beginTime, invokeResult);

        TraceIdUtils.removeTraceId();

        try {
            // rpcStack出栈
            Stack<String> stack = DapThreadContext.getContext(TraceConstants.RPC_STACK);
            MDC.put(Constants.RPC_ID, stack.pop());
        } catch (Throwable e) {
        }

        // 出栈之后将本服务父节点变成当前父节点
        DapThreadContext.putContext(TraceConstants.RPC_PARENT_CLASS_NAME, parentClassName);
    }


    /**
     * 管理和传递rpcid
     * 当前进入切面的类名称
     */
    private String manageRpcIdBeforeMethodInvoked(String classname) {
        try {
            RpcTreeNode tree = DapThreadContext.getContext(TraceConstants.RPC_TREE);
            Stack<String> stack = DapThreadContext.getContext(TraceConstants.RPC_STACK);
            String rpcId, nodeId;

            // 如果线程对象中没有rpcTree和rpcStack，则创建对象
            if (tree == null || stack == null || stack.isEmpty()) {
                rpcId = DapThreadContext.getContext(Constants.RPC_ID);
                if (StringUtils.isEmpty(rpcId))
                    rpcId = LongStringIdUtil.generateId();

                nodeId = classname + rpcId;
                tree = new RpcTreeNode(nodeId, null);
                DapThreadContext.putContext(TraceConstants.RPC_TREE, tree);

                stack = new Stack<>();
                stack.push(rpcId);
                DapThreadContext.putContext(TraceConstants.RPC_STACK, stack);
            } else {
                // 获取最新rpcId
                try {
                    rpcId = stack.peek();
                } catch (Throwable e) {
                    rpcId = LongStringIdUtil.generateId();
                }

                // 获取调用父类
                String parentClassName = DapThreadContext.getContext(TraceConstants.RPC_PARENT_CLASS_NAME);
                // 寻找父节点
                RpcTreeNode parent = tree.findNodeById(parentClassName);
                // 更新rpcid
                if (parent.getChildList() == null || parent.getChildList().isEmpty()) {
                    rpcId = rpcId + ".1";
                } else {
                    rpcId = rpcId + "." + (parent.getChildList().size() + 1);
                }
                try {
                    stack.push(rpcId);
                } catch (Throwable e) {
                }

                // 创建当前节点
                nodeId = classname + rpcId;
                RpcTreeNode node = new RpcTreeNode(nodeId, parentClassName);
                parent.addChild(node);
            }

            DapThreadContext.putContext(TraceConstants.RPC_PARENT_CLASS_NAME, nodeId);
            return rpcId;
        } catch (Exception e) {
            return null;
        }
    }
}
